!> author: 左志华
!> date: 2022-05-26
!>
!> 速度 verlet 积分法 / 蛙跳法
module verlet_m

    implicit none
    private

    public :: velocity_verlet, leap_frog

    abstract interface
        !> `fcn_velocity_verlet` 计算 ODE 的导数
        pure subroutine fcn_velocity_verlet(y, a)
            real, intent(in) :: y(:)   !! 位置
            real, intent(out) :: a(:)  !! 加速度
        end subroutine fcn_velocity_verlet
        !> `fcn_leap_frog` 计算 ODE 的导数
        pure subroutine fcn_leap_frog(y, v, a)
            real, intent(in) :: y(:)   !! 位置
            real, intent(in) :: v(:)   !! 速度
            real, intent(out) :: a(:)  !! 加速度
        end subroutine fcn_leap_frog
    end interface

contains

    !> 速度 verlet 积分法
    !> @note
    !> 1. 在 CFD 中，往往基于边界条件、初始条件驱动粒子初始运动，然后位置、速度通过 verlet 算法得到更新，
    !> 加速度则由新的边界条件、初始条件确定，不幸的是，在大多数 CFD 场景中，
    !> 加速度（力）往往包括与速度相关的粘性项，这在速度 verlet 算法是不适应的，对应只能选用蛙跳法。<br>
    !> 2. 在视频游戏中，在追求**互动性**的需求中，verlet 算法可以通过位置的变化，快速反求出速度。<br>
    !> 3. 在分子动力学中，（好像）加速度仅收势函数影响，也就是与位置相关，这非常适合使用速度 verlet 算法。
    pure subroutine velocity_verlet(func, tspan, y0, v0, t, y, v, a)
        procedure(fcn_velocity_verlet) :: func
        real, intent(in) :: tspan(2)    !! 自变量范围
        real, intent(in) :: y0(:)       !! 位置初值
        real, intent(in) :: v0(:)       !! 速度初值
        real, intent(out) :: t(:)       !! 自变量
        real, intent(out) :: y(:, :)    !! 位置
        real, intent(out) :: v(:, :)    !! 速度
        real, intent(out) :: a(:, :)    !! 加速度
        real :: h
        integer :: i

        h = (tspan(2) - tspan(1))/(size(t) - 1)
        t(1) = tspan(1)
        y(1, :) = y0(:)
        call func(y(1, :), a(1, :))
        v(1, :) = v0(:)

        do i = 1, size(t) - 1
            y(i + 1, :) = y(i, :) + v(i, :)*h + a(i, :)*h**2/2
            call func(y(i + 1, :), a(i + 1, :))     ! @note 计算加速度（力），加速度仅与位置相关
            v(i + 1, :) = v(i, :) + (a(i + 1, :) + a(i, :))*h/2
            t(i + 1) = t(i) + h
        end do

    end subroutine velocity_verlet

    !> 蛙跳法
    !> @note
    !> 1. verlet算法的位移和速度是没有关系的，而蛙跳算法位移和速度有关！<br>
    !> 2. 速度比位移、加速度慢半个步长。
    !> @todo
    !> 与龙格库塔法进行性能测试，测试相同工作量计算效率
    pure subroutine leap_frog(func, tspan, y0, v0, t, y, v, a, sync_v)
        procedure(fcn_leap_frog) :: func
        real, intent(in) :: tspan(2)    !! 自变量范围
        real, intent(in) :: y0(:)       !! 位置初值
        real, intent(in) :: v0(:)       !! 速度初值
        real, intent(out) :: t(:)       !! 自变量
        real, intent(out) :: y(:, :)    !! 位置
        real, intent(out) :: v(:, :)    !! 速度
        real, intent(out) :: a(:, :)    !! 加速度
        logical, intent(in), optional :: sync_v     !! 是否同步速度，默认为否
                                                    !! T: 同步，位置、速度、加速度同步更新
                                                    !! F: 不同步，速度比位置、加速度慢半步长
        real :: h, vtmp(size(y, 2))
        integer :: i
        logical :: sync

        if (present(sync_v)) then
            sync = sync_v
        else
            sync = .false.
        end if

        ! 第 0 步
        h = (tspan(2) - tspan(1))/(size(t) - 1)
        t(1) = tspan(1)
        v(1, :) = v0(:)
        y(1, :) = y0(:)
        call func(y(1, :), v(1, :), a(1, :))

        ! 计算第 1 步位置、速度、加速度
        ! @note
        ! 1. 速度推进半个步长；<br>
        ! 2. 位置、加速度推进一个步长。
        t(2) = t(1) + h
        vtmp = v(1, :) + a(1, :)*h/2    ! @note: 认为半步长内加速度是值不变
        y(2, :) = y(1, :) + vtmp*h
        call func(y(2, :), vtmp + a(1, :)*h/2, a(2, :))

        if (sync) then
            v(2, :) = vtmp + a(1, :)*h/2
        else
            v(2, :) = vtmp
        end if

        ! 之后步
        do i = 2, size(t) - 1
            t(i + 1) = t(i) + h
            vtmp = vtmp + a(i, :)*h
            y(i + 1, :) = y(i, :) + vtmp*h
            ! @note 速度使用加速度多推进了半个步长，同步位置和速度，计算加速度。
            call func(y(i + 1, :), vtmp + a(i, :)*h/2, a(i + 1, :))
            if (sync) then
                v(i + 1, :) = vtmp + a(i, :)*h/2  ! @note 采用更新的加速度推进半个步长
            else
                v(i + 1, :) = vtmp
            end if
        end do

    end subroutine leap_frog

end module verlet_m
